// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/serializers/compact_row_serializer.h>
#include <pollux/serializers/row_serializer.h>

#include <melon/lang/bits.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/row/compact_row.h>

namespace kumo::pollux::serializer {
    namespace {
        using TRowSize = uint32_t;

        class CompactRowVectorSerializer : public RowSerializer<row::CompactRow> {
        public:
            explicit CompactRowVectorSerializer(
                memory::MemoryPool *pool,
                const VectorSerde::Options *options)
                : RowSerializer<row::CompactRow>(pool, options) {
            }

        private:
            void serializeRanges(
                const row::CompactRow &row,
                const melon::Range<const IndexRange *> &ranges,
                char *rawBuffer,
                const std::vector<vector_size_t> &rowSize) override {
                size_t offset = 0;
                vector_size_t index = 0;
                for (const auto &range: ranges) {
                    if (range.size == 1) {
                        // Fast path for single-row serialization.
                        *reinterpret_cast<TRowSize *>(rawBuffer + offset) =
                                melon::Endian::big(rowSize[index]);
                        auto size =
                                row.serialize(range.begin, rawBuffer + offset + sizeof(TRowSize));
                        offset += size + sizeof(TRowSize);
                        ++index;
                    } else {
                        raw_vector<size_t> offsets(range.size);
                        for (auto i = 0; i < range.size; ++i, ++index) {
                            // Write raw size. Needs to be in big endian order.
                            *(TRowSize *) (rawBuffer + offset) = melon::Endian::big(rowSize[index]);
                            offsets[i] = offset + sizeof(TRowSize);
                            offset += rowSize[index] + sizeof(TRowSize);
                        }
                        // Write row data for all rows in range.
                        row.serialize(range.begin, range.size, offsets.data(), rawBuffer);
                    }
                }
            }
        };
    } // namespace

    void CompactRowVectorSerde::estimateSerializedSize(
        const row::CompactRow *compactRow,
        const melon::Range<const vector_size_t *> &rows,
        vector_size_t **sizes) {
        compactRow->serializedRowSizes(rows, sizes);
    }

    std::unique_ptr<IterativeVectorSerializer>
    CompactRowVectorSerde::createIterativeSerializer(
        RowTypePtr /* type */,
        int32_t /* numRows */,
        StreamArena *streamArena,
        const Options *options) {
        return std::make_unique<CompactRowVectorSerializer>(
            streamArena->pool(), options);
    }

    void CompactRowVectorSerde::deserialize(
        ByteInputStream *source,
        pollux::memory::MemoryPool *pool,
        RowTypePtr type,
        RowVectorPtr *result,
        const Options *options) {
        std::vector<std::string_view> serializedRows;
        std::vector<std::unique_ptr<std::string> > serializedBuffers;
        RowDeserializer<std::string_view>::deserialize<RowIteratorImpl>(
            source, serializedRows, serializedBuffers, options);

        if (serializedRows.empty()) {
            *result = BaseVector::create<RowVector>(type, 0, pool);
            return;
        }

        *result = row::CompactRow::deserialize(serializedRows, type, pool);
    }

    // static
    void CompactRowVectorSerde::registerVectorSerde() {
        pollux::registerVectorSerde(std::make_unique<CompactRowVectorSerde>());
    }

    // static
    void CompactRowVectorSerde::registerNamedVectorSerde() {
        pollux::registerNamedVectorSerde(
            VectorSerde::Kind::kCompactRow,
            std::make_unique<CompactRowVectorSerde>());
    }
} // namespace kumo::pollux::serializer
